import torch
import torch.nn as nn
import sys
import os

cur_path = os.path.split(__file__)[0]
sys.path.append(cur_path)

from typing import List
from BuildingBlocks import (
    Encoder,
    Decoder,
    FinalConv,
    DoubleConv,
    ExtResNetBlock,
    SingleConv,
)
from utils.normalization import Normalizer


def create_feature_maps(init_channel_number, number_of_fmaps):
    return [init_channel_number * 2**k for k in range(number_of_fmaps)]


class UNet3D(nn.Module):
    """
    3DUnet model from
    `"3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation"
        <https://arxiv.org/pdf/1606.06650.pdf>`.
    Args:
        in_channels (int): number of input channels
        out_channels (int): number of output segmentation masks;
            Note that that the of out_channels might correspond to either
            different semantic classes or to different binary segmentation mask.
            It's up to the user of the class to interpret the out_channels and
            use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
            or BCEWithLogitsLoss (two-class) respectively)
        f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
            of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
        final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the
            final 1x1 convolution, otherwise apply nn.Softmax. MUST be True if nn.BCELoss (two-class) is used
            to train the model. MUST be False if nn.CrossEntropyLoss (multi-class) is used to train the model.
        layer_order (string): determines the order of layers
            in `SingleConv` module. e.g. 'crg' stands for Conv3d+ReLU+GroupNorm3d.
            See `SingleConv` for more info
        init_channel_number (int): number of feature maps in the first conv layer of the encoder; default: 64
        num_groups (int): number of groups for the GroupNorm
    """

    def __init__(
        self,
        in_channels,
        nlayers=4,
        f_maps=16,
        layer_order="crg",
        num_groups=8,
        **kwargs
    ):
        super(UNet3D, self).__init__()

        if isinstance(f_maps, int):
            # use 4 levels in the encoder path as suggested in the paper
            f_maps = create_feature_maps(f_maps, number_of_fmaps=nlayers)

        # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
        # uses DoubleConv as a basic_module for the Encoder
        encoders = []
        for i, out_feature_num in enumerate(f_maps):
            if i == 0:
                encoder = Encoder(
                    in_channels,
                    out_feature_num,
                    apply_pooling=False,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            else:
                encoder = Encoder(
                    f_maps[i - 1],
                    out_feature_num,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            encoders.append(encoder)

        self.encoders = nn.ModuleList(encoders)

        # create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
        # uses DoubleConv as a basic_module for the Decoder
        decoders = []
        reversed_f_maps = list(reversed(f_maps))
        for i in range(len(reversed_f_maps) - 1):
            in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
            out_feature_num = reversed_f_maps[i + 1]
            decoder = Decoder(
                in_feature_num,
                out_feature_num,
                basic_module=DoubleConv,
                conv_layer_order=layer_order,
                num_groups=num_groups,
            )
            decoders.append(decoder)

        self.decoders = nn.ModuleList(decoders)

    def forward(self, x):
        # encoder part
        encoders_features = []
        for encoder in self.encoders:
            x = encoder(x)
            # reverse the encoder outputs to be aligned with the decoder
            encoders_features.insert(0, x)

        # remove the last encoder's output from the list
        # !!remember: it's the 1st in the list

        encoders_features = encoders_features[1:]
        # decoder part
        for decoder, encoder_features in zip(self.decoders, encoders_features):
            # pass the output from the corresponding encoder and the output
            # of the previous decoder
            x = decoder(encoder_features, x)

        # apply final_activation (i.e. Sigmoid or Softmax) only for prediction. During training the network outputs
        # logits and it's up to the user to normalize it before visualising with tensorboard or computing validation metric

        return x


class UNet3DWithSamplePoints(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        hidden_channels: int,
        num_levels: int,
        use_position_input: bool = True,
        params=None,
    ):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        self.unet3d = UNet3D(
            in_channels=in_channels, nlayers=num_levels, f_maps=hidden_channels
        )
        self.out_ln1 = nn.LayerNorm(hidden_channels)
        # self.out_ln2 = nn.Sequential(
        #     nn.Linear(hidden_channels, hidden_channels),
        #     nn.GELU(),
        #     nn.LayerNorm(hidden_channels),
        # )

        self.hidden_channels = hidden_channels
        try:
            satistic_times = params.dataset_size // params.batch_size
        except:
            satistic_times = 500

        self.u_normlizer = Normalizer(in_channels, satistic_times, filter_zero=False)

        self.global_pool = nn.AdaptiveAvgPool3d(1)

        # self.fc = nn.Sequential(
        #     nn.Linear(hidden_channels, hidden_channels * 2),
        #     nn.GELU(),
        #     nn.Linear(hidden_channels * 2, 1),
        # )

    def voxel_expand(self, voxel_left):  # B C Z X Y
        indices = torch.arange(31, -1, -1, dtype=torch.long, device=voxel_left.device)
        voxel_right = voxel_left[:, :, :, indices, :]
        voxel_merge = torch.cat((voxel_left, voxel_right), axis=3)
        return voxel_merge

    def forward(self, x, output_points, half=False):
        # Use super class to get the output of the UNet3D
        # 0, 1, 2, 3, 4
        B, C, H, W, D = x.shape

        x = (
            self.u_normlizer(x.permute(0, 2, 3, 4, 1).reshape(-1, C))
            .reshape(B, H, W, D, C)
            .permute(0, 4, 1, 2, 3)
        )

        x = self.unet3d.forward(x)

        # sample_points to 5D by adding two dimensions
        if half:
            x = self.voxel_expand(x)

        if isinstance(output_points, List):
            rt_x = []
            for idx, output_point in enumerate(output_points):

                cur_x = (
                    torch.nn.functional.grid_sample(
                        x[idx : idx + 1], output_point, align_corners=False
                    )
                    .squeeze()
                    .T
                ).view(-1, self.hidden_channels)

                rt_x.append(cur_x)

            rt_x = torch.cat(rt_x, dim=0)
        else:
            rt_x = (
                torch.nn.functional.grid_sample(x, output_points, align_corners=False)
                .squeeze()
                .T
            ).view(-1, self.hidden_channels)

        sampled_u_feature = self.out_ln1(rt_x)
        unsampled_u_feature = x.permute(0, 2, 3, 4, 1).reshape(B, -1, self.hidden_channels)
        # pred_cd_hidden = self.global_pool(x).squeeze().view(-1,self.hidden_channels)
        pred_cd_hidden = None

        return sampled_u_feature, unsampled_u_feature

        # return self.out_ln(rt_x)


class Deep3DConv(nn.Module):
    def __init__(
        self,
        in_channels,
        nlayers=4,
        f_maps=16,
        layer_order="crg",
        num_groups=8,
        **kwargs
    ):
        super(Deep3DConv, self).__init__()

        if isinstance(f_maps, int):
            # use 4 levels in the encoder path as suggested in the paper
            f_maps = create_feature_maps(f_maps, number_of_fmaps=nlayers)

        # create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
        # uses DoubleConv as a basic_module for the Encoder
        encoders = []
        for i, out_feature_num in enumerate(f_maps):
            if i == 0:
                encoder = Encoder(
                    in_channels,
                    out_feature_num,
                    apply_pooling=False,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            else:
                encoder = Encoder(
                    f_maps[i - 1],
                    out_feature_num,
                    basic_module=DoubleConv,
                    conv_layer_order=layer_order,
                    num_groups=num_groups,
                )
            encoders.append(encoder)

        self.encoders = nn.ModuleList(encoders)

    def forward(self, x):
        for encoder in self.encoders:
            x = encoder(x)
        return x


if __name__ == "__main__":
    # net = UNet3DWithSamplePoints(1, 64, 64, 4, True)
    # sdf = torch.randn([7, 1, 64, 64, 64])
    # outputpoints = torch.randn([7, 1471, 1, 1, 3])
    # output = net.forward(sdf, outputpoints, half=False)
    # print(output.shape)

    net = Deep3DConv(in_channels=1, nlayers=7, f_maps=32).cuda()
    sdf = torch.randn([7, 1, 64, 64, 64]).cuda()
    output = net.forward(sdf)
    print(output.shape)
